{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "28df5c07-95ab-409d-9304-d22333f3780d",
   "metadata": {},
   "source": [
    "# Extraction and Analysis over a Fidelity Multi-Fund Annual Report\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_cloud_services/blob/main/examples/extract/asset_manager_fund_analysis.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "In this notebook we show you how to create an agentic document workflow over a complex document that contains annual reports for multiple funds - each fund reports financials in a standardized reporting structure, and it's all consolidated in the same document.\n",
    "\n",
    "We show you how to create a workflow that parses the document, splits it into subsections per fund, and runs extraction per subsection. The extracted results can be compiled into a single CSV file.\n",
    "\n",
    "![](asset_manager_fund_analysis.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cda2e5e9-fe9d-42d9-9387-f529d970ff7b",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "\n",
    "### Data\n",
    "\n",
    "We load the consolidated document from Fidelity's Shareholder Report portal.\n",
    "\n",
    "**To get the original Doc**\n",
    "\n",
    "Link: https://fundresearch.fidelity.com/prospectus/sec\n",
    "- Go to row \"Fidelity Advisor Asset Manager® 20% Fund Class A\", and click \"Annual\" under \"Financial Statements & Other Information\"\n",
    "\n",
    "**Or, just download it below (we're hosting it on Dropbox)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d563afc1-ccfa-410a-97a6-85041f8551f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://www.dropbox.com/scl/fi/bhrtivs7b2gz3yhrr4t4s/fidelity_fund.pdf?rlkey=ha2loufvuer1c07u47k68hgji&st=ev66x31t&dl=1\" -O data/asset_manager_fund_analysis/fidelity_fund.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "098d71d4-dd70-4fd8-879e-557ba2d78899",
   "metadata": {},
   "source": [
    "### Everything else\n",
    "\n",
    "In the below sections we define the required environment variables and LLM/embedding model.\n",
    "\n",
    "Sign up for a LlamaCloud account here: https://cloud.llamaindex.ai/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b12c54b-5a0f-4c9a-86da-155aa549bb3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# API access to llama-cloud\n",
    "os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"llx-...\"\n",
    "\n",
    "# Using OpenAI API for embeddings/llms\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2b6741-26ce-4472-8be7-00dd037d9ea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set LlamaCloud project/organization ID\n",
    "# project_id = None\n",
    "# organization_id = None\n",
    "project_id = \"2fef999e-1073-40e6-aeb3-1f3c0e64d99b\"\n",
    "organization_id = \"43b88c8f-e488-46f6-9013-698e3d2e374a\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f61c14b-9813-4886-8869-4d8d973d04c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "# set LLM, embedding model\n",
    "embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")\n",
    "llm = OpenAI(model=\"gpt-4.1\")\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d12de2e-87f3-4c43-af26-06913e1b41cb",
   "metadata": {},
   "source": [
    "## Create Splitter\n",
    "\n",
    "The document is a consolidated regulatory filing that contains multiple sub-entities - the 7 different Asset Manager funds - with identical reporting structures.\n",
    "\n",
    "Our first task is to create a document splitter that can identify the specific page numbers corresponding to each fund. This is a pre-requisite for the downstream task of doing per-fund extraction and consolidation.\n",
    "\n",
    "The below block uses LlamaParse to parse the full document and output the result in per-page Markdown nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46cc4516-022a-4299-ae51-c5525882ae9b",
   "metadata": {},
   "source": [
    "#### Setup Parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0408e7b1-ce23-4422-b389-8f69f6e37a6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Retrying llama_cloud_services.parse.utils.make_api_request.<locals>._make_request in 4.0 seconds as it raised RemoteProtocolError: Server disconnected without sending a response..\n",
      "Retrying llama_cloud_services.parse.utils.make_api_request.<locals>._make_request in 4.0 seconds as it raised RemoteProtocolError: Server disconnected without sending a response..\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id 36c97d60-995c-4b82-824b-3553c55d93e3\n"
     ]
    }
   ],
   "source": [
    "from llama_cloud_services import LlamaParse\n",
    "\n",
    "parser = LlamaParse(\n",
    "    premium_mode=True,\n",
    "    result_type=\"markdown\",\n",
    "    project_id=project_id,\n",
    "    organization_id=organization_id,\n",
    ")\n",
    "result = await parser.aparse(\"./data/asset_manager_fund_analysis/fidelity_fund.pdf\")\n",
    "markdown_nodes = await result.aget_markdown_nodes(split_by_page=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4bb2c88-d4ca-4eb8-bcc5-420ddf5bc0a0",
   "metadata": {},
   "source": [
    "### Define Splitting Functions\n",
    "\n",
    "We then define a set of functions to find the \"splits\". \n",
    "1. We first identify the split categories given a user description. (e.g. what are the precise labels we want to split on?) \n",
    "2. We then split the document according to the split categories. \n",
    "\n",
    "On (2), we do this by using LLMs on each page to detect if a split starts on that page. We then count all pages from the start of the current split to the start on the next split as within the current split."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3610297-9319-4b31-a351-8c1458fbb08c",
   "metadata": {},
   "source": [
    "#### Find Split Categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "951ee2b9-5d90-479b-9e92-247037c66742",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from typing import List, Optional, Dict\n",
    "from llama_index.core.schema import TextNode\n",
    "from llama_index.core.llms import LLM\n",
    "from llama_index.core.async_utils import run_jobs, asyncio_run\n",
    "from llama_index.core.prompts import ChatPromptTemplate, ChatMessage\n",
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "split_category_prompt = \"\"\"\\\n",
    "You are an AI document assistant tasked with finding the 'split categories' given a user description and the document text.\n",
    "- The split categories is a list of string tags from the document that correspond to the user description.\n",
    "- Do not make up split categories. \n",
    "- Do not include category tags that don't fit the user description,\\\n",
    "for instance subcategories or extraneous titles.\n",
    "- Do not exclude category tags that do fit the user description. \n",
    "\n",
    "For instance, if the user asks to \"find all top-level sections of an ArXiv paper\", then a sample output would be:\n",
    "[\"1. Introduction\", \"2. Related Work\", \"3. Methodology\", \"4. Experiments\", \"5. Conclusion\"]\n",
    "\n",
    "The split description and document text are given below. \n",
    "\n",
    "Split description:\n",
    "{split_description}\n",
    "\n",
    "Here is the document text:\n",
    "{document_text}\n",
    "    \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class SplitCategories(BaseModel):\n",
    "    \"\"\"A list of all split categories from a document.\"\"\"\n",
    "\n",
    "    split_categories: List[str]\n",
    "\n",
    "\n",
    "async def afind_split_categories(\n",
    "    split_description: str,\n",
    "    nodes: List[TextNode],\n",
    "    llm: Optional[LLM] = None,\n",
    "    page_limit: Optional[int] = 5,\n",
    ") -> List[str]:\n",
    "    \"\"\"Find split categories given a user description and the page limit.\n",
    "    \n",
    "    These categories will then be used to find the exact splits of the document. \n",
    "    \n",
    "    NOTE: with the page limit there is an assumption that the categories are found in the first few pages,\\\n",
    "    for instance in the table of contents. This does not account for the case where the categories are \\\n",
    "    found throughout the document. \n",
    "    \n",
    "    \"\"\"\n",
    "    llm = llm or OpenAI(model=\"gpt-4.1\")\n",
    "\n",
    "    chat_template = ChatPromptTemplate(\n",
    "        [\n",
    "            ChatMessage.from_str(split_category_prompt, \"user\"),\n",
    "        ]\n",
    "    )\n",
    "    nodes_head = nodes[:page_limit] if page_limit is not None else nodes\n",
    "    doc_text = \"\\n-----\\n\".join(\n",
    "        [n.get_content(metadata_mode=\"all\") for n in nodes_head]\n",
    "    )\n",
    "\n",
    "    result = await llm.astructured_predict(\n",
    "        SplitCategories,\n",
    "        chat_template,\n",
    "        split_description=split_description,\n",
    "        document_text=doc_text,\n",
    "    )\n",
    "    return result.split_categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccd7f58d-f48c-4b87-96ca-baf132d71a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "split_categories = await afind_split_categories(\n",
    "    \"Find and split by the main funds in this document, should be listed in the first few pages\",\n",
    "    markdown_nodes,\n",
    "    llm=llm,\n",
    "    page_limit=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7047d15-48d3-4469-91d3-037529a47140",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Fidelity Asset Manager® 20%',\n",
       " 'Fidelity Asset Manager® 30%',\n",
       " 'Fidelity Asset Manager® 40%',\n",
       " 'Fidelity Asset Manager® 50%',\n",
       " 'Fidelity Asset Manager® 60%',\n",
       " 'Fidelity Asset Manager® 70%',\n",
       " 'Fidelity Asset Manager® 85%']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_categories"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7fd8746-4a45-4f79-afca-30d28b99a183",
   "metadata": {},
   "source": [
    "#### Split Document based on Existing Categories/Rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471cf952-713f-47eb-a40e-c2318564fa19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# go through each node (on a page level), and tag node\n",
    "\n",
    "\n",
    "split_prompt = \"\"\"\\\n",
    "You are an AI document assistant tasked with extracting out splits from a document text according to a certain set of rules. \n",
    "\n",
    "You are given a chunk of the document text at a time. \n",
    "You are responsible for determining if the chunk of the document text corresponds to the beginning of a split. \n",
    "\n",
    "We've listed general rules below, and the user has also provided their own rules to find a split. Please extract\n",
    "out the splits according to the defined schema. \n",
    "\n",
    "General Rules: \n",
    "- You should ONLY extract out a split if the document text contains the beginning of a split.\n",
    "- If the document text contains the beginning of two or more splits (e.g. there are multiple sections on a single page), then \\\n",
    "return all splits in the output.\n",
    "- If the text does not correspond to the beginning of any split, then return a blank list. \n",
    "- A valid split must be clearly delineated in the document text according to the user rules. \\\n",
    "Do NOT identify a split if it is mentioned, but is not actually the start of a split in the document text.\n",
    "- If you do find one or more splits, please output the split_name according to the format \\\"{split_key}_X\\\", \\\n",
    "where X is a short tag corresponding to the split. \n",
    "\n",
    "Split key:\n",
    "{split_key}\n",
    "\n",
    "User-defined rules:\n",
    "{split_rules}\n",
    "\n",
    "\n",
    "Here is the chunk text:\n",
    "{chunk_text}\n",
    "    \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class SplitOutput(BaseModel):\n",
    "    \"\"\"The metadata for a given split start given a chunk.\"\"\"\n",
    "\n",
    "    split_name: str = Field(\n",
    "        ..., description=\"The name of the split (in the format \\{split_key\\}_X)\"\n",
    "    )\n",
    "    split_description: str = Field(\n",
    "        ..., description=\"A short description corresponding to the split.\"\n",
    "    )\n",
    "    page_number: int = Field(..., description=\"Page number of the split.\")\n",
    "\n",
    "\n",
    "class SplitsOutput(BaseModel):\n",
    "    \"\"\"A list of all splits given a chunk.\"\"\"\n",
    "\n",
    "    splits: List[SplitOutput]\n",
    "\n",
    "\n",
    "async def atag_splits_in_node(\n",
    "    split_rules: str, split_key: str, node: TextNode, llm: Optional[LLM] = None\n",
    "):\n",
    "    \"\"\"Tag split in a single node.\"\"\"\n",
    "    llm = llm or OpenAI(model=\"gpt-4o\")\n",
    "\n",
    "    chat_template = ChatPromptTemplate(\n",
    "        [\n",
    "            ChatMessage.from_str(split_prompt, \"user\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    result = await llm.astructured_predict(\n",
    "        SplitsOutput,\n",
    "        chat_template,\n",
    "        split_rules=split_rules,\n",
    "        split_key=split_key,\n",
    "        chunk_text=node.get_content(metadata_mode=\"all\"),\n",
    "    )\n",
    "    return result.splits\n",
    "\n",
    "\n",
    "async def afind_splits(\n",
    "    split_rules: str, split_key: str, nodes: List[TextNode], llm: Optional[LLM] = None\n",
    ") -> Dict:\n",
    "    \"\"\"Find splits.\"\"\"\n",
    "\n",
    "    # tag each node with split or no-split\n",
    "    tasks = [atag_splits_in_node(split_rules, split_key, n, llm=llm) for n in nodes]\n",
    "    async_results = await run_jobs(tasks, workers=8, show_progress=True)\n",
    "    all_splits = [s for r in async_results for s in r]\n",
    "\n",
    "    split_name_to_pages = defaultdict(list)\n",
    "\n",
    "    split_idx = 0\n",
    "    for idx, n in enumerate(nodes):\n",
    "        cur_page = n.metadata[\"page_number\"]\n",
    "\n",
    "        # update the current split if needed\n",
    "        while (\n",
    "            split_idx + 1 < len(all_splits)\n",
    "            and all_splits[split_idx + 1].page_number <= cur_page\n",
    "        ):\n",
    "            split_idx += 1\n",
    "\n",
    "        # add page number to the current split\n",
    "        if all_splits[split_idx].page_number <= cur_page:\n",
    "            split_name = all_splits[split_idx].split_name\n",
    "            split_name_to_pages[split_name].append(cur_page)\n",
    "\n",
    "    # print(all_splits)\n",
    "    # print(split_name_to_pages)\n",
    "    # raise Exception\n",
    "\n",
    "    return split_name_to_pages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "737573c6-771a-4fde-bcb3-221b5e3616f7",
   "metadata": {},
   "source": [
    "#### Find Categories and then Run Splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5552735-e4f7-4a76-a58a-dad2a5bfab3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# put it all together - detect categories, then split document based on those categories\n",
    "async def afind_categories_and_splits(\n",
    "    split_description: str,\n",
    "    split_key: str,\n",
    "    nodes: List[TextNode],\n",
    "    additional_split_rules: Optional[str] = None,\n",
    "    llm: Optional[LLM] = None,\n",
    "    page_limit: int = 5,\n",
    "    verbose: bool = False,\n",
    "):\n",
    "    \"\"\"Find categories and then splits.\"\"\"\n",
    "    categories = await afind_split_categories(\n",
    "        split_description, nodes, llm=llm, page_limit=page_limit\n",
    "    )\n",
    "    if verbose:\n",
    "        print(f\"Split categories: {categories}\")\n",
    "    full_split_rules = f\"\"\"Please split by these categories: {categories}\"\"\"\n",
    "    if additional_split_rules:\n",
    "        full_split_rules += f\"\\n\\n\\n{additional_split_rules}\"\n",
    "\n",
    "    return await afind_splits(full_split_rules, split_key, nodes, llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "279ad80d-3c1d-4a73-9c61-a9bc819571be",
   "metadata": {},
   "source": [
    "#### Define document-specific split rules\n",
    "\n",
    "Now that we have the general splitter, let's define the rules to help parse this Fidelity document.\n",
    "\n",
    "**NOTE**: Ideally the valid names are automatically identified, through a separate pass in the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ea71c0-ee30-47a1-88d2-62899cd57d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fidelity_split_rules = \"\"\"\n",
    "# - You must split by the name of the fund\n",
    "# - Here are the valid names - do not include subcategories:\n",
    "#     Fidelity Asset Manager® 20%\n",
    "#     Fidelity Asset Manager® 30%\n",
    "#     Fidelity Asset Manager® 40%\n",
    "#     Fidelity Asset Manager® 50%\n",
    "#     Fidelity Asset Manager® 60%\n",
    "#     Fidelity Asset Manager® 70%\n",
    "#     Fidelity Asset Manager® 85%\n",
    "# - Each fund will have a list of tables underneath it, like schedule of investments, financial statements\n",
    "# - Each fund usually has schedule of investments right underneath it\n",
    "# - Do not tag the cover page/table of contents\n",
    "# \"\"\"\n",
    "# fidelity_split_key = \"fidelity_asset_manager\"\n",
    "\n",
    "\n",
    "fidelity_split_description = \"Find and split by the main funds in this document, should be listed in the first few pages\"\n",
    "fidelity_split_rules = \"\"\"\n",
    "- You must split by the name of the fund\n",
    "- Each fund will have a list of tables underneath it, like schedule of investments, financial statements\n",
    "- Each fund usually has schedule of investments right underneath it \n",
    "- Do not tag the cover page/table of contents\n",
    "\"\"\"\n",
    "fidelity_split_key = \"fidelity_asset_manager\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e41c96-4ccf-43ff-b24d-1bc3a4828825",
   "metadata": {},
   "source": [
    "#### Try it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a344810c-acf3-49db-9dc7-519b55f371be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split categories: ['Fidelity Asset Manager® 20%', 'Fidelity Asset Manager® 30%', 'Fidelity Asset Manager® 40%', 'Fidelity Asset Manager® 50%', 'Fidelity Asset Manager® 60%', 'Fidelity Asset Manager® 70%', 'Fidelity Asset Manager® 85%']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████| 120/120 [00:28<00:00,  4.20it/s]\n"
     ]
    }
   ],
   "source": [
    "split_name_to_pages = await afind_categories_and_splits(\n",
    "    fidelity_split_description,\n",
    "    fidelity_split_key,\n",
    "    markdown_nodes,\n",
    "    additional_split_rules=fidelity_split_rules,\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20304fd-031d-4556-bfa0-0c626175df76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "defaultdict(list,\n",
       "            {'fidelity_asset_manager_20': [2,\n",
       "              3,\n",
       "              4,\n",
       "              5,\n",
       "              6,\n",
       "              7,\n",
       "              8,\n",
       "              9,\n",
       "              10,\n",
       "              11,\n",
       "              12,\n",
       "              13,\n",
       "              93,\n",
       "              94,\n",
       "              108,\n",
       "              109,\n",
       "              110,\n",
       "              111],\n",
       "             'fidelity_asset_manager_30': [14,\n",
       "              15,\n",
       "              16,\n",
       "              17,\n",
       "              18,\n",
       "              19,\n",
       "              20,\n",
       "              21,\n",
       "              22,\n",
       "              23,\n",
       "              24,\n",
       "              25,\n",
       "              99,\n",
       "              101,\n",
       "              112],\n",
       "             'fidelity_asset_manager_40': [26,\n",
       "              27,\n",
       "              28,\n",
       "              29,\n",
       "              30,\n",
       "              31,\n",
       "              32,\n",
       "              33,\n",
       "              34,\n",
       "              35,\n",
       "              36,\n",
       "              88,\n",
       "              102,\n",
       "              106,\n",
       "              113],\n",
       "             'fidelity_asset_manager_50': [37,\n",
       "              38,\n",
       "              39,\n",
       "              40,\n",
       "              41,\n",
       "              42,\n",
       "              43,\n",
       "              44,\n",
       "              45,\n",
       "              46,\n",
       "              47,\n",
       "              91,\n",
       "              114],\n",
       "             'fidelity_asset_manager_60': [48,\n",
       "              49,\n",
       "              50,\n",
       "              51,\n",
       "              52,\n",
       "              53,\n",
       "              54,\n",
       "              55,\n",
       "              56,\n",
       "              57,\n",
       "              58,\n",
       "              103],\n",
       "             'fidelity_asset_manager_70': [59,\n",
       "              60,\n",
       "              61,\n",
       "              62,\n",
       "              63,\n",
       "              64,\n",
       "              65,\n",
       "              66,\n",
       "              67,\n",
       "              68,\n",
       "              69,\n",
       "              70,\n",
       "              96],\n",
       "             'fidelity_asset_manager_85': [71,\n",
       "              72,\n",
       "              73,\n",
       "              74,\n",
       "              75,\n",
       "              76,\n",
       "              77,\n",
       "              78,\n",
       "              79,\n",
       "              80,\n",
       "              81,\n",
       "              82,\n",
       "              83,\n",
       "              84,\n",
       "              85,\n",
       "              86,\n",
       "              87,\n",
       "              89,\n",
       "              90,\n",
       "              92,\n",
       "              95,\n",
       "              97,\n",
       "              98,\n",
       "              100,\n",
       "              104,\n",
       "              105,\n",
       "              107,\n",
       "              115,\n",
       "              116,\n",
       "              117,\n",
       "              118,\n",
       "              119,\n",
       "              120]})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_name_to_pages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32092f81-0af0-4240-80d7-73e72a44170c",
   "metadata": {},
   "source": [
    "## Run Per-Fund Extraction, Consolidate Results\n",
    "\n",
    "We then use LlamaExtract on the aggregated text for each split (page numbers corresponding to each fund), to extract out a Pydantic `FundData` object for each fund.\n",
    "\n",
    "We then consolidate them together into a CSV."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9dc62cc-122e-4bec-a1e9-aa7f594be167",
   "metadata": {},
   "source": [
    "#### Define Fund Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb4fdd4-a8ed-4c90-a0ff-7112a8019c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define output schema\n",
    "from pydantic import BaseModel\n",
    "from typing import Optional\n",
    "from datetime import date\n",
    "\n",
    "\n",
    "class FundData(BaseModel):\n",
    "    \"\"\"Concise fund data extraction schema optimized for LLM extraction\"\"\"\n",
    "\n",
    "    # Identifiers\n",
    "    fund_name: str = Field(\n",
    "        ...,\n",
    "        description=\"Full fund name exactly as it appears, e.g. 'Fidelity Asset Manager® 20%'\",\n",
    "    )\n",
    "    target_equity_pct: Optional[int] = Field(\n",
    "        None,\n",
    "        description=\"Target equity percentage from fund name (20, 30, 40, 50, 60, 70, or 85)\",\n",
    "    )\n",
    "    report_date: Optional[str] = Field(\n",
    "        None, description=\"Report date in YYYY-MM-DD format, e.g. '2024-09-30'\"\n",
    "    )\n",
    "\n",
    "    # Asset Allocation (as percentages, e.g. 27.4 for 27.4%)\n",
    "    equity_pct: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Actual equity allocation percentage from 'Equity Central Funds' section\",\n",
    "    )\n",
    "    fixed_income_pct: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Fixed income allocation percentage from 'Fixed-Income Central Funds' section\",\n",
    "    )\n",
    "    money_market_pct: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Money market allocation percentage from 'Money Market Central Funds' section\",\n",
    "    )\n",
    "    other_pct: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Other investments percentage (Treasury + Investment Companies + other)\",\n",
    "    )\n",
    "\n",
    "    # Primary Share Class Metrics (use the main retail class, usually named after the fund)\n",
    "    nav: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Net Asset Value per share for the main retail class (e.g. Asset Manager 20% class)\",\n",
    "    )\n",
    "    net_assets_usd: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Total net assets in USD for the main retail class from 'Net Asset Value' section\",\n",
    "    )\n",
    "    expense_ratio: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Expense ratio as percentage (e.g. 0.48 for 0.48%) from Financial Highlights\",\n",
    "    )\n",
    "    management_fee: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Management fee rate as percentage from Financial Highlights or Notes\",\n",
    "    )\n",
    "\n",
    "    # Performance (as percentages)\n",
    "    one_year_return: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"One-year total return percentage from Financial Highlights (e.g. 13.74 for 13.74%)\",\n",
    "    )\n",
    "    portfolio_turnover: Optional[float] = Field(\n",
    "        None, description=\"Portfolio turnover rate percentage from Financial Highlights\"\n",
    "    )\n",
    "\n",
    "    # Risk Metrics (in USD)\n",
    "    equity_futures_notional: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Net notional amount of equity futures contracts (positive if net long, negative if net short)\",\n",
    "    )\n",
    "    bond_futures_notional: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Net notional amount of bond/treasury futures contracts (positive if net long, negative if net short)\",\n",
    "    )\n",
    "\n",
    "    # Fund Flows (in USD)\n",
    "    net_investment_income: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Net investment income for the period from Statement of Operations\",\n",
    "    )\n",
    "    total_distributions: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Total distributions to shareholders from Statement of Changes in Net Assets\",\n",
    "    )\n",
    "    net_asset_change: Optional[float] = Field(\n",
    "        None,\n",
    "        description=\"Net change in assets from beginning to end of period (end minus beginning net assets)\",\n",
    "    )\n",
    "\n",
    "\n",
    "class FundComparisonData(BaseModel):\n",
    "    \"\"\"Flattened data optimized for CSV export and analysis\"\"\"\n",
    "\n",
    "    funds: list[FundData]\n",
    "\n",
    "    def to_csv_rows(self) -> list[dict]:\n",
    "        \"\"\"Convert to list of dictionaries for CSV export\"\"\"\n",
    "        return [fund.dict() for fund in self.funds]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b681ae85-fc76-45c6-a191-c30b84e748f5",
   "metadata": {},
   "source": [
    "#### Setup LlamaExtract "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b113b682-a823-4ad4-921e-0f2444da0a04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "deleting existing agent\n"
     ]
    }
   ],
   "source": [
    "from llama_cloud_services import LlamaExtract\n",
    "from llama_cloud.core.api_error import ApiError\n",
    "from llama_cloud import ExtractConfig\n",
    "\n",
    "\n",
    "# Optionally, add your project id/organization id\n",
    "llama_extract = LlamaExtract(\n",
    "    show_progress=True,\n",
    "    check_interval=5,\n",
    "    project_id=project_id,\n",
    "    organization_id=organization_id,\n",
    ")\n",
    "\n",
    "\n",
    "try:\n",
    "    existing_agent = llama_extract.get_agent(name=\"FundDataExtractor2\")\n",
    "    if existing_agent:\n",
    "        print(f\"deleting existing agent\")\n",
    "        # Deletion can take some time since all underlying files will be purged\n",
    "        llama_extract.delete_agent(existing_agent.id)\n",
    "except ApiError as e:\n",
    "    if e.status_code == 404:\n",
    "        pass\n",
    "    else:\n",
    "        raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f4ce337-e8b3-48c2-b25e-ed141b503735",
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_config = ExtractConfig(\n",
    "    extraction_mode=\"BALANCED\",\n",
    ")\n",
    "\n",
    "extract_agent = llama_extract.create_agent(\n",
    "    \"FundDataExtractor2\", data_schema=FundData, config=extract_config\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63c6b41-dcc4-4354-ba9d-be462931b45b",
   "metadata": {},
   "source": [
    "#### Define Extraction Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e37b7050-924d-4afe-9dd3-c77f11d9d698",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cloud_services.extract import SourceText\n",
    "from typing import List, Optional, Dict\n",
    "\n",
    "\n",
    "async def aextract_data_over_split(\n",
    "    split_name: str,\n",
    "    page_numbers: List[int],\n",
    "    nodes: List[TextNode],\n",
    "    llm: Optional[LLM] = None,\n",
    ") -> FundData:\n",
    "    \"\"\"Extract fund data for a given split.\"\"\"\n",
    "\n",
    "    # combine node text that matches the page numbers\n",
    "    filtered_nodes = [n for n in nodes if n.metadata[\"page_number\"] in page_numbers]\n",
    "    filtered_text = \"\\n-------\\n\".join(\n",
    "        [n.get_content(metadata_mode=\"all\") for n in filtered_nodes]\n",
    "    )\n",
    "    result_dict = (\n",
    "        await extract_agent.aextract(SourceText(text_content=filtered_text))\n",
    "    ).data\n",
    "\n",
    "    fund_data = FundData.model_validate(result_dict)\n",
    "\n",
    "    return fund_data\n",
    "\n",
    "\n",
    "async def aextract_data_over_splits(\n",
    "    split_name_to_pages: Dict[str, List],\n",
    "    nodes: List[TextNode],\n",
    "    llm: Optional[LLM] = None,\n",
    "):\n",
    "    \"\"\"Extract fund data for each split, aggregate.\"\"\"\n",
    "    tasks = [\n",
    "        aextract_data_over_split(split_name, page_numbers, nodes, llm=llm)\n",
    "        for split_name, page_numbers in split_name_to_pages.items()\n",
    "    ]\n",
    "    all_fund_data = await run_jobs(tasks, workers=8, show_progress=True)\n",
    "    return FundComparisonData(funds=all_fund_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89b768db-877d-460f-8f32-61ba954a492f",
   "metadata": {},
   "source": [
    "#### Try it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5be6348-2533-4513-9d01-05b8cca5f5e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_fund_data = await aextract_data_over_splits(\n",
    "    split_name_to_pages, markdown_nodes, llm=llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121df331-0b35-48ec-b056-686e5d56fa45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fund_name</th>\n",
       "      <th>target_equity_pct</th>\n",
       "      <th>report_date</th>\n",
       "      <th>equity_pct</th>\n",
       "      <th>fixed_income_pct</th>\n",
       "      <th>money_market_pct</th>\n",
       "      <th>other_pct</th>\n",
       "      <th>nav</th>\n",
       "      <th>net_assets_usd</th>\n",
       "      <th>expense_ratio</th>\n",
       "      <th>management_fee</th>\n",
       "      <th>one_year_return</th>\n",
       "      <th>portfolio_turnover</th>\n",
       "      <th>equity_futures_notional</th>\n",
       "      <th>bond_futures_notional</th>\n",
       "      <th>net_investment_income</th>\n",
       "      <th>total_distributions</th>\n",
       "      <th>net_asset_change</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Fidelity Asset Manager® 20%</td>\n",
       "      <td>20</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>27.4</td>\n",
       "      <td>48.6</td>\n",
       "      <td>16.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>13.95</td>\n",
       "      <td>3.170882e+09</td>\n",
       "      <td>0.48</td>\n",
       "      <td>0.43</td>\n",
       "      <td>13.74</td>\n",
       "      <td>24.0</td>\n",
       "      <td>-308348.0</td>\n",
       "      <td>159650906.0</td>\n",
       "      <td>180771003.0</td>\n",
       "      <td>1.815011e+08</td>\n",
       "      <td>148593411.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Fidelity Asset Manager® 30%</td>\n",
       "      <td>30</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>37.4</td>\n",
       "      <td>48.4</td>\n",
       "      <td>6.3</td>\n",
       "      <td>8.1</td>\n",
       "      <td>12.21</td>\n",
       "      <td>1.375227e+09</td>\n",
       "      <td>0.49</td>\n",
       "      <td>0.44</td>\n",
       "      <td>16.20</td>\n",
       "      <td>19.0</td>\n",
       "      <td>-117123430.0</td>\n",
       "      <td>60683344.0</td>\n",
       "      <td>60219520.0</td>\n",
       "      <td>6.074714e+07</td>\n",
       "      <td>61425310.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Fidelity Asset Manager® 40%</td>\n",
       "      <td>40</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>47.3</td>\n",
       "      <td>43.3</td>\n",
       "      <td>6.4</td>\n",
       "      <td>3.0</td>\n",
       "      <td>13.40</td>\n",
       "      <td>1.544767e+09</td>\n",
       "      <td>0.51</td>\n",
       "      <td>NaN</td>\n",
       "      <td>18.24</td>\n",
       "      <td>13.0</td>\n",
       "      <td>-131196445.0</td>\n",
       "      <td>68340188.0</td>\n",
       "      <td>59239637.0</td>\n",
       "      <td>3.794509e+07</td>\n",
       "      <td>117313232.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Fidelity Asset Manager® 50%</td>\n",
       "      <td>50</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>57.1</td>\n",
       "      <td>37.3</td>\n",
       "      <td>2.1</td>\n",
       "      <td>3.5</td>\n",
       "      <td>21.44</td>\n",
       "      <td>7.407487e+09</td>\n",
       "      <td>0.56</td>\n",
       "      <td>0.58</td>\n",
       "      <td>20.37</td>\n",
       "      <td>12.0</td>\n",
       "      <td>-550922128.0</td>\n",
       "      <td>288445875.0</td>\n",
       "      <td>223428532.0</td>\n",
       "      <td>2.197314e+08</td>\n",
       "      <td>771033500.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Fidelity Asset Manager® 60%</td>\n",
       "      <td>60</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>63.8</td>\n",
       "      <td>32.2</td>\n",
       "      <td>0.7</td>\n",
       "      <td>3.3</td>\n",
       "      <td>16.32</td>\n",
       "      <td>2.282371e+09</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.54</td>\n",
       "      <td>22.31</td>\n",
       "      <td>17.0</td>\n",
       "      <td>-103026215.0</td>\n",
       "      <td>127423594.0</td>\n",
       "      <td>84848944.0</td>\n",
       "      <td>7.662887e+07</td>\n",
       "      <td>630046421.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Fidelity Asset Manager® 70%</td>\n",
       "      <td>70</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>73.7</td>\n",
       "      <td>22.3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>3.3</td>\n",
       "      <td>29.09</td>\n",
       "      <td>4.273012e+09</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.54</td>\n",
       "      <td>24.06</td>\n",
       "      <td>14.0</td>\n",
       "      <td>-162028290.0</td>\n",
       "      <td>202620656.0</td>\n",
       "      <td>119678101.0</td>\n",
       "      <td>1.121974e+08</td>\n",
       "      <td>935150295.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Fidelity Asset Manager® 85%</td>\n",
       "      <td>85</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>84.6</td>\n",
       "      <td>11.2</td>\n",
       "      <td>1.9</td>\n",
       "      <td>3.3</td>\n",
       "      <td>27.85</td>\n",
       "      <td>4.394805e+09</td>\n",
       "      <td>0.94</td>\n",
       "      <td>0.63</td>\n",
       "      <td>27.23</td>\n",
       "      <td>16.0</td>\n",
       "      <td>105407845.0</td>\n",
       "      <td>129202345.0</td>\n",
       "      <td>60293654.0</td>\n",
       "      <td>3.695516e+08</td>\n",
       "      <td>935099874.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                     fund_name  target_equity_pct report_date  equity_pct  \\\n",
       "0  Fidelity Asset Manager® 20%                 20  2024-09-30        27.4   \n",
       "1  Fidelity Asset Manager® 30%                 30  2024-09-30        37.4   \n",
       "2  Fidelity Asset Manager® 40%                 40  2024-09-30        47.3   \n",
       "3  Fidelity Asset Manager® 50%                 50  2024-09-30        57.1   \n",
       "4  Fidelity Asset Manager® 60%                 60  2024-09-30        63.8   \n",
       "5  Fidelity Asset Manager® 70%                 70  2024-09-30        73.7   \n",
       "6  Fidelity Asset Manager® 85%                 85  2024-09-30        84.6   \n",
       "\n",
       "   fixed_income_pct  money_market_pct  other_pct    nav  net_assets_usd  \\\n",
       "0              48.6              16.0        8.0  13.95    3.170882e+09   \n",
       "1              48.4               6.3        8.1  12.21    1.375227e+09   \n",
       "2              43.3               6.4        3.0  13.40    1.544767e+09   \n",
       "3              37.3               2.1        3.5  21.44    7.407487e+09   \n",
       "4              32.2               0.7        3.3  16.32    2.282371e+09   \n",
       "5              22.3               1.0        3.3  29.09    4.273012e+09   \n",
       "6              11.2               1.9        3.3  27.85    4.394805e+09   \n",
       "\n",
       "   expense_ratio  management_fee  one_year_return  portfolio_turnover  \\\n",
       "0           0.48            0.43            13.74                24.0   \n",
       "1           0.49            0.44            16.20                19.0   \n",
       "2           0.51             NaN            18.24                13.0   \n",
       "3           0.56            0.58            20.37                12.0   \n",
       "4           0.64            0.54            22.31                17.0   \n",
       "5           0.64            0.54            24.06                14.0   \n",
       "6           0.94            0.63            27.23                16.0   \n",
       "\n",
       "   equity_futures_notional  bond_futures_notional  net_investment_income  \\\n",
       "0                -308348.0            159650906.0            180771003.0   \n",
       "1             -117123430.0             60683344.0             60219520.0   \n",
       "2             -131196445.0             68340188.0             59239637.0   \n",
       "3             -550922128.0            288445875.0            223428532.0   \n",
       "4             -103026215.0            127423594.0             84848944.0   \n",
       "5             -162028290.0            202620656.0            119678101.0   \n",
       "6              105407845.0            129202345.0             60293654.0   \n",
       "\n",
       "   total_distributions  net_asset_change  \n",
       "0         1.815011e+08       148593411.0  \n",
       "1         6.074714e+07        61425310.0  \n",
       "2         3.794509e+07       117313232.0  \n",
       "3         2.197314e+08       771033500.0  \n",
       "4         7.662887e+07       630046421.0  \n",
       "5         1.121974e+08       935150295.0  \n",
       "6         3.695516e+08       935099874.0  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "all_fund_data_df = pd.DataFrame(all_fund_data.to_csv_rows())\n",
    "all_fund_data_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589fcb4d-39a8-4aef-81ca-3702f095a261",
   "metadata": {},
   "source": [
    "## Define Full Workflow\n",
    "\n",
    "We put everything together into a LlamaIndex workflow! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66fae1ce-5372-4c63-813b-0def195feb7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Event,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Context,\n",
    "    Workflow,\n",
    "    step,\n",
    ")\n",
    "\n",
    "\n",
    "class ParseDocEvent(Event):\n",
    "    nodes: List[TextNode]\n",
    "\n",
    "\n",
    "class DocSplitEvent(Event):\n",
    "    split_name_to_pages: Dict[str, List[int]]\n",
    "    nodes: List[TextNode]\n",
    "\n",
    "\n",
    "class FidelityFundExtraction(Workflow):\n",
    "    \"\"\"\n",
    "    Workflow to extract data from a solar panel datasheet and generate a comparison report\n",
    "    against provided design requirements.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        parser: LlamaParse,\n",
    "        extract_agent: LlamaExtract,\n",
    "        split_description: str = fidelity_split_description,\n",
    "        split_rules: str = fidelity_split_rules,\n",
    "        split_key: str = fidelity_split_key,\n",
    "        llm: Optional[LLM] = None,\n",
    "        **kwargs\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.parser = parser\n",
    "        self.extract_agent = extract_agent\n",
    "        self.split_description = split_description\n",
    "        self.split_rules = split_rules\n",
    "        self.split_key = split_key\n",
    "        self.llm = llm\n",
    "\n",
    "    @step\n",
    "    async def parse_doc(self, ctx: Context, ev: StartEvent) -> ParseDocEvent:\n",
    "        \"\"\"Parse document into markdown nodes.\"\"\"\n",
    "        result = await parser.aparse(file_path=ev.file_path)\n",
    "        markdown_nodes = await result.aget_markdown_nodes(split_by_page=True)\n",
    "        return ParseDocEvent(nodes=markdown_nodes)\n",
    "\n",
    "    @step\n",
    "    async def find_splits(self, ctx: Context, ev: ParseDocEvent) -> DocSplitEvent:\n",
    "        split_name_to_pages = await afind_categories_and_splits(\n",
    "            self.split_description,\n",
    "            self.split_key,\n",
    "            ev.nodes,\n",
    "            additional_split_rules=self.split_rules,\n",
    "            llm=llm,\n",
    "            verbose=True,\n",
    "        )\n",
    "        return DocSplitEvent(\n",
    "            split_name_to_pages=split_name_to_pages,\n",
    "            nodes=ev.nodes,\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def run_extraction(self, ctx: Context, ev: DocSplitEvent) -> StopEvent:\n",
    "        all_fund_data = await aextract_data_over_splits(\n",
    "            ev.split_name_to_pages, ev.nodes, llm=self.llm\n",
    "        )\n",
    "        all_fund_data_df = pd.DataFrame(all_fund_data.to_csv_rows())\n",
    "        return StopEvent(\n",
    "            result={\n",
    "                \"all_fund_data\": all_fund_data,\n",
    "                \"all_fund_data_df\": all_fund_data_df,\n",
    "            }\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "870368c6-a164-4f84-973f-d7534fce28cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = FidelityFundExtraction(\n",
    "    parser=parser, extract_agent=extract_agent, verbose=True, timeout=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0813641-105a-4321-b46a-27a2c48df39c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fidelity_fund_extraction.html\n"
     ]
    }
   ],
   "source": [
    "from llama_index.utils.workflow import draw_all_possible_flows\n",
    "\n",
    "draw_all_possible_flows(\n",
    "    FidelityFundExtraction,\n",
    "    filename=\"fidelity_fund_extraction.html\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0f14ac3-c681-46a9-8664-0fc8bd1513f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = await workflow.run(\n",
    "    file_path=\"./data/asset_manager_fund_analysis/fidelity_fund.pdf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a4608b-0807-4d3e-9166-d393d21fb970",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_fund_data_df = result[\"all_fund_data_df\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe931879-3fbc-4c26-95ed-a7b6790a0a5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fund_name</th>\n",
       "      <th>target_equity_pct</th>\n",
       "      <th>report_date</th>\n",
       "      <th>equity_pct</th>\n",
       "      <th>fixed_income_pct</th>\n",
       "      <th>money_market_pct</th>\n",
       "      <th>other_pct</th>\n",
       "      <th>nav</th>\n",
       "      <th>net_assets_usd</th>\n",
       "      <th>expense_ratio</th>\n",
       "      <th>management_fee</th>\n",
       "      <th>one_year_return</th>\n",
       "      <th>portfolio_turnover</th>\n",
       "      <th>equity_futures_notional</th>\n",
       "      <th>bond_futures_notional</th>\n",
       "      <th>net_investment_income</th>\n",
       "      <th>total_distributions</th>\n",
       "      <th>net_asset_change</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Fidelity Asset Manager® 20%</td>\n",
       "      <td>20</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>27.4</td>\n",
       "      <td>48.6</td>\n",
       "      <td>16.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>13.95</td>\n",
       "      <td>3.170882e+09</td>\n",
       "      <td>0.48</td>\n",
       "      <td>0.48</td>\n",
       "      <td>13.74</td>\n",
       "      <td>24.0</td>\n",
       "      <td>-30834897.0</td>\n",
       "      <td>159650906.0</td>\n",
       "      <td>180771003.0</td>\n",
       "      <td>181501111.0</td>\n",
       "      <td>148593411.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Fidelity Asset Manager® 30%</td>\n",
       "      <td>30</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>37.4</td>\n",
       "      <td>48.4</td>\n",
       "      <td>6.3</td>\n",
       "      <td>8.1</td>\n",
       "      <td>12.21</td>\n",
       "      <td>1.375227e+09</td>\n",
       "      <td>0.49</td>\n",
       "      <td>0.44</td>\n",
       "      <td>16.20</td>\n",
       "      <td>19.0</td>\n",
       "      <td>-117123430.0</td>\n",
       "      <td>60683344.0</td>\n",
       "      <td>60219520.0</td>\n",
       "      <td>60747136.0</td>\n",
       "      <td>61325310.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Fidelity Asset Manager® 40%</td>\n",
       "      <td>40</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>47.3</td>\n",
       "      <td>43.3</td>\n",
       "      <td>6.4</td>\n",
       "      <td>3.0</td>\n",
       "      <td>13.40</td>\n",
       "      <td>1.544767e+09</td>\n",
       "      <td>0.51</td>\n",
       "      <td>0.43</td>\n",
       "      <td>18.24</td>\n",
       "      <td>13.0</td>\n",
       "      <td>-131196445.0</td>\n",
       "      <td>68340188.0</td>\n",
       "      <td>59239637.0</td>\n",
       "      <td>-37925052.0</td>\n",
       "      <td>117313232.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Fidelity Asset Manager® 50%</td>\n",
       "      <td>50</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>57.1</td>\n",
       "      <td>37.3</td>\n",
       "      <td>2.1</td>\n",
       "      <td>3.5</td>\n",
       "      <td>21.44</td>\n",
       "      <td>7.407487e+09</td>\n",
       "      <td>0.56</td>\n",
       "      <td>0.53</td>\n",
       "      <td>20.37</td>\n",
       "      <td>12.0</td>\n",
       "      <td>-550922128.0</td>\n",
       "      <td>288445875.0</td>\n",
       "      <td>223428532.0</td>\n",
       "      <td>219731373.0</td>\n",
       "      <td>771033500.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Fidelity Asset Manager® 60%</td>\n",
       "      <td>60</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>63.8</td>\n",
       "      <td>32.2</td>\n",
       "      <td>0.7</td>\n",
       "      <td>3.3</td>\n",
       "      <td>16.32</td>\n",
       "      <td>2.282371e+09</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.54</td>\n",
       "      <td>22.31</td>\n",
       "      <td>17.0</td>\n",
       "      <td>-103026215.0</td>\n",
       "      <td>127423594.0</td>\n",
       "      <td>84848944.0</td>\n",
       "      <td>76628870.0</td>\n",
       "      <td>630046421.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Fidelity Asset Manager® 70%</td>\n",
       "      <td>70</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>73.7</td>\n",
       "      <td>22.3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>3.3</td>\n",
       "      <td>29.09</td>\n",
       "      <td>4.273012e+09</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.16</td>\n",
       "      <td>24.06</td>\n",
       "      <td>14.0</td>\n",
       "      <td>-162028290.0</td>\n",
       "      <td>202620656.0</td>\n",
       "      <td>119678101.0</td>\n",
       "      <td>112197413.0</td>\n",
       "      <td>935150295.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Fidelity Asset Manager® 85%</td>\n",
       "      <td>85</td>\n",
       "      <td>2024-09-30</td>\n",
       "      <td>84.6</td>\n",
       "      <td>11.2</td>\n",
       "      <td>1.9</td>\n",
       "      <td>3.3</td>\n",
       "      <td>27.85</td>\n",
       "      <td>4.394805e+09</td>\n",
       "      <td>0.94</td>\n",
       "      <td>0.52</td>\n",
       "      <td>27.23</td>\n",
       "      <td>16.0</td>\n",
       "      <td>104412803.0</td>\n",
       "      <td>129202345.0</td>\n",
       "      <td>60293654.0</td>\n",
       "      <td>51390318.0</td>\n",
       "      <td>935099874.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                     fund_name  target_equity_pct report_date  equity_pct  \\\n",
       "0  Fidelity Asset Manager® 20%                 20  2024-09-30        27.4   \n",
       "1  Fidelity Asset Manager® 30%                 30  2024-09-30        37.4   \n",
       "2  Fidelity Asset Manager® 40%                 40  2024-09-30        47.3   \n",
       "3  Fidelity Asset Manager® 50%                 50  2024-09-30        57.1   \n",
       "4  Fidelity Asset Manager® 60%                 60  2024-09-30        63.8   \n",
       "5  Fidelity Asset Manager® 70%                 70  2024-09-30        73.7   \n",
       "6  Fidelity Asset Manager® 85%                 85  2024-09-30        84.6   \n",
       "\n",
       "   fixed_income_pct  money_market_pct  other_pct    nav  net_assets_usd  \\\n",
       "0              48.6              16.0        8.0  13.95    3.170882e+09   \n",
       "1              48.4               6.3        8.1  12.21    1.375227e+09   \n",
       "2              43.3               6.4        3.0  13.40    1.544767e+09   \n",
       "3              37.3               2.1        3.5  21.44    7.407487e+09   \n",
       "4              32.2               0.7        3.3  16.32    2.282371e+09   \n",
       "5              22.3               1.0        3.3  29.09    4.273012e+09   \n",
       "6              11.2               1.9        3.3  27.85    4.394805e+09   \n",
       "\n",
       "   expense_ratio  management_fee  one_year_return  portfolio_turnover  \\\n",
       "0           0.48            0.48            13.74                24.0   \n",
       "1           0.49            0.44            16.20                19.0   \n",
       "2           0.51            0.43            18.24                13.0   \n",
       "3           0.56            0.53            20.37                12.0   \n",
       "4           0.64            0.54            22.31                17.0   \n",
       "5           0.64            0.16            24.06                14.0   \n",
       "6           0.94            0.52            27.23                16.0   \n",
       "\n",
       "   equity_futures_notional  bond_futures_notional  net_investment_income  \\\n",
       "0              -30834897.0            159650906.0            180771003.0   \n",
       "1             -117123430.0             60683344.0             60219520.0   \n",
       "2             -131196445.0             68340188.0             59239637.0   \n",
       "3             -550922128.0            288445875.0            223428532.0   \n",
       "4             -103026215.0            127423594.0             84848944.0   \n",
       "5             -162028290.0            202620656.0            119678101.0   \n",
       "6              104412803.0            129202345.0             60293654.0   \n",
       "\n",
       "   total_distributions  net_asset_change  \n",
       "0          181501111.0       148593411.0  \n",
       "1           60747136.0        61325310.0  \n",
       "2          -37925052.0       117313232.0  \n",
       "3          219731373.0       771033500.0  \n",
       "4           76628870.0       630046421.0  \n",
       "5          112197413.0       935150295.0  \n",
       "6           51390318.0       935099874.0  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_fund_data_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd06ac2-c012-496d-ba00-acddb08dadfc",
   "metadata": {},
   "source": [
    "## Run Analysis over Compiled Data\n",
    "\n",
    "Now that we've created an aggregated dataframe from this report, we can now do good ol' Pandas analysis or use LLMs to help do text-to-CSV over it! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4de5633-090e-436c-869f-a1f8015c28d9",
   "metadata": {},
   "source": [
    "#### Standard Pandas analysis\n",
    "\n",
    "The conservative 20% fund is actually the most return-efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bf74cee-3b4b-42b4-a6df-6567287dea82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    0.501460\n",
       "1    0.433155\n",
       "2    0.385624\n",
       "3    0.356743\n",
       "4    0.349687\n",
       "5    0.326459\n",
       "6    0.321868\n",
       "Name: return_per_risk, dtype: float64"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calculate return per unit of risk (equity allocation)\n",
    "all_fund_data_df[\"return_per_risk\"] = (\n",
    "    all_fund_data_df[\"one_year_return\"] / all_fund_data_df[\"equity_pct\"]\n",
    ")\n",
    "all_fund_data_df[\"return_per_risk\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b097e10b-af9e-4d90-b587-14495cdab94f",
   "metadata": {},
   "source": [
    "The lower equity allocation funds consistently run higher than their target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48373eb3-5ebe-4610-b05d-3f0239925d97",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    7.4\n",
       "1    7.4\n",
       "2    7.3\n",
       "3    7.1\n",
       "4    3.8\n",
       "5    3.7\n",
       "6   -0.4\n",
       "Name: drift, dtype: float64"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# How far do actual allocations drift from targets?\n",
    "all_fund_data_df[\"drift\"] = (\n",
    "    all_fund_data_df[\"equity_pct\"] - all_fund_data_df[\"target_equity_pct\"]\n",
    ")\n",
    "all_fund_data_df[\"drift\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76580f49-f5d9-4a60-b5d2-89160f16272e",
   "metadata": {},
   "source": [
    "#### Text-to-Pandas Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e99a64d-b112-439c-871c-fe7905839d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.experimental.query_engine import PandasQueryEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6975fead-7278-4ee0-880a-5e65843e6c3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd_query_engine = PandasQueryEngine(\n",
    "    df=all_fund_data_df, verbose=True, synthesize_response=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ae8b4d-ce69-455e-8e2b-e6b277ff632a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Pandas Instructions:\n",
      "```\n",
      "df[df['equity_pct'] > df['target_equity_pct'] + 5]\n",
      "```\n",
      "> Pandas Output:                      fund_name  target_equity_pct report_date  equity_pct  \\\n",
      "0  Fidelity Asset Manager® 20%                 20  2024-09-30        27.4   \n",
      "1  Fidelity Asset Manager® 30%                 30  2024-09-30        37.4   \n",
      "2  Fidelity Asset Manager® 40%                 40  2024-09-30        47.3   \n",
      "3  Fidelity Asset Manager® 50%                 50  2024-09-30        57.1   \n",
      "\n",
      "   fixed_income_pct  money_market_pct  other_pct    nav  net_assets_usd  \\\n",
      "0              48.6              16.0        8.0  13.95    3.170882e+09   \n",
      "1              48.4               6.3        8.1  12.21    1.375227e+09   \n",
      "2              43.3               6.4        3.0  13.40    1.544767e+09   \n",
      "3              37.3               2.1        3.5  21.44    7.407487e+09   \n",
      "\n",
      "   expense_ratio  management_fee  one_year_return  portfolio_turnover  \\\n",
      "0           0.48            0.48            13.74                24.0   \n",
      "1           0.49            0.44            16.20                19.0   \n",
      "2           0.51            0.43            18.24                13.0   \n",
      "3           0.56            0.53            20.37                12.0   \n",
      "\n",
      "   equity_futures_notional  bond_futures_notional  net_investment_income  \\\n",
      "0              -30834897.0            159650906.0            180771003.0   \n",
      "1             -117123430.0             60683344.0             60219520.0   \n",
      "2             -131196445.0             68340188.0             59239637.0   \n",
      "3             -550922128.0            288445875.0            223428532.0   \n",
      "\n",
      "   total_distributions  net_asset_change  return_per_risk  drift  \n",
      "0          181501111.0       148593411.0         0.501460    7.4  \n",
      "1           60747136.0        61325310.0         0.433155    7.4  \n",
      "2          -37925052.0       117313232.0         0.385624    7.3  \n",
      "3          219731373.0       771033500.0         0.356743    7.1  \n",
      "The following funds have an actual equity allocation that is more than 5% higher than their target equity allocation as of the 2024-09-30 report date:\n",
      "\n",
      "1. **Fidelity Asset Manager® 20%**\n",
      "   - Target Equity: 20%\n",
      "   - Actual Equity: 27.4%\n",
      "   - Difference: +7.4%\n",
      "\n",
      "2. **Fidelity Asset Manager® 30%**\n",
      "   - Target Equity: 30%\n",
      "   - Actual Equity: 37.4%\n",
      "   - Difference: +7.4%\n",
      "\n",
      "3. **Fidelity Asset Manager® 40%**\n",
      "   - Target Equity: 40%\n",
      "   - Actual Equity: 47.3%\n",
      "   - Difference: +7.3%\n",
      "\n",
      "4. **Fidelity Asset Manager® 50%**\n",
      "   - Target Equity: 50%\n",
      "   - Actual Equity: 57.1%\n",
      "   - Difference: +7.1%\n",
      "\n",
      "All four funds listed have an actual equity allocation that exceeds their target by more than 5%.\n"
     ]
    }
   ],
   "source": [
    "response = pd_query_engine.query(\n",
    "    \"Show me all funds where the actual equity allocation is more than 5% higher than the target\"\n",
    ")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_parse",
   "language": "python",
   "name": "llama_parse"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
